package com.inspur.edp.ai.flow.engine;

import com.fasterxml.jackson.annotation.JsonIgnore;
import com.inspur.edp.ai.flow.exception.FlowException;
import com.inspur.edp.ai.flow.metadata.AbsNode;
import com.inspur.edp.ai.flow.metadata.Edge;
import lombok.Data;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * @author lizhaorui
 * @date 2025/9/19
 * @description
 */

@Data
public class Flow {

    private String id;

    private String parentId;

    private List<AbsNode> nodes;

    private List<Edge> edges;

    @JsonIgnore
    public List<Edge> getOutputEdges(String sourceNodeId, String sourcePort) {
        if (this.edges == null) {
            return new ArrayList<>();
        } else {
            return this.edges.stream().filter(new Predicate<Edge>() {
                @Override
                public boolean test(Edge edge) {
                    return edge.getSourceNodeId().equals(sourceNodeId) && edge.getSourcePort().equals(sourcePort);
                }
            }).collect(Collectors.toList());
        }
    }

    @JsonIgnore
    public List<Edge> getInputEdges(String targetNodeId) {
        if (this.edges == null) {
            return new ArrayList<>();
        } else {
            return this.edges.stream().filter(new Predicate<Edge>() {
                @Override
                public boolean test(Edge edge) {
                    return edge.getTargetNodeId().equals(targetNodeId);
                }
            }).collect(Collectors.toList());
        }
    }


    @JsonIgnore
    public List<Edge> getOutputEdges(String sourceNodeId) {
        if (this.edges == null) {
            return new ArrayList<>();
        } else {
            return this.edges.stream().filter(new Predicate<Edge>() {
                @Override
                public boolean test(Edge edge) {
                    return edge.getSourceNodeId().equals(sourceNodeId);
                }
            }).collect(Collectors.toList());
        }
    }

    @JsonIgnore
    public AbsNode getHeadNode() {
        if (this.nodes == null) {
            return null;
        } else {
            Map<String, String> nodeIdMap = new HashMap<>();
            for (Edge edge : this.edges) {
                nodeIdMap.put(edge.getTargetNodeId(), edge.getTargetNodeId());
            }
            List<AbsNode> headNodes = new ArrayList<>();
            for (AbsNode node : this.nodes) {
                if (!nodeIdMap.containsKey(node.getId())) {
                    headNodes.add(node);
                }
            }
            if (headNodes.size() > 1) {
                throw new FlowException("has more than 1 head node.");
            }
            if (headNodes.size() == 0) {
                throw new FlowException("not has head node.");
            }
            return headNodes.get(0);
        }
    }


    @JsonIgnore
    public AbsNode getStartNode() {
        if (this.nodes == null) {
            return null;
        } else {
            List<AbsNode> queryNodes = this.nodes.stream().filter(new Predicate<AbsNode>() {
                @Override
                public boolean test(AbsNode node) {
                    return node.getKind().equals("start");
                }
            }).collect(Collectors.toList());
            if (queryNodes.size() > 0) {
                return queryNodes.get(0);
            } else {
                return null;
            }
        }
    }

    @JsonIgnore
    public List<String> getEndNodeIds() {
        if (this.nodes == null) {
            return null;
        } else {
            Map<String, String> hasTargets = new HashMap<>();
            if (this.edges != null) {
                for (Edge edge : this.edges) {
                    hasTargets.put(edge.getSourceNodeId(), edge.getSourceNodeId());
                }
            }
            List<String> endNodeIds = new ArrayList<>();
            for (AbsNode node : this.nodes) {
                if (!hasTargets.containsKey(node.getId())) {
                    endNodeIds.add(node.getId());
                }
            }
            return endNodeIds;
        }
    }

    @JsonIgnore
    public AbsNode getNodeById(String id) {
        if (this.nodes == null) {
            return null;
        } else {
            List<AbsNode> queryNodes = this.nodes.stream().filter(new Predicate<AbsNode>() {
                @Override
                public boolean test(AbsNode node) {
                    return node.getId().equals(id);
                }
            }).collect(Collectors.toList());
            if (queryNodes.size() > 0) {
                return queryNodes.get(0);
            } else {
                return null;
            }
        }
    }
}